#coding=utf-8

from plate_utils import *

class GenPlate:
    def __init__(self, font_ch = font_ch_plate, font_eng = font_en_plate, plate_bg_dir = plate_bg_dir):
        self.fontC =  ImageFont.truetype(font_ch, 43, 0)
        self.fontE =  ImageFont.truetype(font_eng, 60, 0)
        self.img = np.array(Image.new("RGB", (226,70),(255,255,255)))
        self.bg  = cv2.resize(cv2.imread(plate_template_img_path),(226,70))
        self.smu = cv2.imread(plate_noise_img_path)
        self.noplates_path = []
        for parent,parent_folder,filenames in os.walk(plate_bg_dir):
            for filename in filenames:
                path = parent+"/"+filename
                self.noplates_path.append(path)


    def draw(self,val):
        offset= 2

        self.img[0:70,offset+8:offset+8+23]= GenCh(self.fontC,val[0])
        self.img[0:70,offset+8+23+6:offset+8+23+6+23]= GenCh1(self.fontE,val[1])
        for i in range(5):
            base = offset+8+23+6+23+17 +i*23 + i*6
            self.img[0:70, base  : base+23]= GenCh1(self.fontE,val[i+2])
        return self.img

    def generate(self,text):
        if len(text) == licence_plate_chars_count:
            # fg = self.draw(text.decode(encoding="utf-8"))
            fg = self.draw(text)
            fg = cv2.bitwise_not(fg)
            com = cv2.bitwise_or(fg,self.bg)
            # com = rot(com,rand_int(60)-30,30)
            com = rot(com, rand_int(40) - 20, 20)
            com = rotRandrom(com, 15, 8 )
            com = append_noise(com,self.smu)

            # com = t_factor(com)
            com = random_envirment(com,self.noplates_path)
            com = add_blur(com, 1 + rand_int(2))
            com = add_noise(com)
            return com

    @staticmethod
    def gen_plate_string():
        plateStr = ''
        for cpos in range(licence_plate_chars_count):
            if cpos == 0:
                plateStr += PLATE_CHARS_TOTAL[rand_int(31)]
            elif cpos == 1:
                plateStr += PLATE_CHARS_TOTAL[41 + rand_int(24)]
            else:
                plateStr += PLATE_CHARS_TOTAL[31 + rand_int(34)]

        return plateStr

    def gen_one(self,size = (plate_width_default,plate_height_default), channels = 1):
        text = self.gen_plate_string()
        img = self.generate(text)
        img = cv2.resize(img, size)

        if channels == 1:
            img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

        return text, img

    def gen_batch(self, batchSize, outputPath, size):
        if not os.path.exists(outputPath):
            os.makedirs(outputPath)
        for i in xrange(batchSize):
            text , img = self.gen_one(size)
            filename = os.path.join(outputPath, str(i).zfill(5) + '_' + text + ".jpg")
            cv2.imwrite(filename, img)

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--font_ch', default = font_ch_plate)
    parser.add_argument('--font_en', default = font_en_plate)
    parser.add_argument('--bg_dir', default = plate_bg_dir )
    parser.add_argument('--out_dir', default = 'plate_tf/plate_train', help = 'output dir')
    parser.add_argument('--make_num', default = 20, type = int, help = 'num')
    parser.add_argument('--img_w', default= plate_width_default, type = int, help='num')
    parser.add_argument('--img_h', default= plate_height_default, type = int, help='num')
    return parser.parse_args()

def main(args):
    g = GenPlate(args.font_ch, args.font_en, args.bg_dir)
    g.gen_batch(args.make_num, args.out_dir, (args.img_w, args.img_h))

if __name__ == '__main__':
    main(parse_args())
